#!/usr/bin/env python3
# map/map_loader.py
# 地图加载和管理：加载 yaml/pgm、坐标转换、障碍应用
import os
import yaml
import cv2
import numpy as np
from math import dist
from PyQt5.QtGui import QPixmap, QImage
import pickle

import logging

logger = logging.getLogger(__name__)


class MapManager:
    def __init__(self, logger=None):
        if logger is None:
            logger = logging.getLogger(__name__)
        self.logger = logger
        self.map_image = None  # cv2 BGR
        self.map_pixmap = None  # QPixmap
        self.map_width = 0
        self.map_height = 0
        self.resolution = 0.05
        self.origin = [0.0, 0.0]
        self.node_coordinates = []
        self.adj_matrix = np.array([])
        self.node_num = 0
        self.obstacle_layer = np.zeros((0, 0), dtype=np.uint8)
        self.original_node_coordinates = []
        self.original_adj_matrix = np.array([])
        self.original_node_num = 0
        self.current_obstacle_layer_index = -1
        self.obstacle_layer_file_num = 0

        self.map_dir_path = "../datas"
        self.load_map()
        self.update_obstacle_layer_file_num()
        self.load_current_obstacle_layer()
        self.load_graph_data()
        self.apply_obstacles()

    def update_obstacle_layer_file_num(self):
        root = os.path.join(os.path.dirname(__file__), self.map_dir_path)
        preset_files = [
            f
            for f in os.listdir(root)
            if f.startswith("obstacle_layer")
            and f.endswith(".pkl")
            and f != "obstacle_layer.pkl"
        ]
        logger.info(f"检测到预设障碍层文件: {preset_files}")
        if preset_files:
            indices = [
                int(f[len("obstacle_layer") : -len(".pkl")]) for f in preset_files
            ]  # obstacle_layer1.pkl -> 1
            self.obstacle_layer_file_num = max(indices)
        else:
            self.obstacle_layer_file_num = 0
        self.logger.info(f"检测到 {self.obstacle_layer_file_num} 个预设障碍层文件")

    def load_current_obstacle_layer(self):
        """加载当前障碍层（默认或预设）"""
        root = os.path.join(os.path.dirname(__file__), self.map_dir_path)
        if self.current_obstacle_layer_index == -1:
            path = os.path.join(root, "obstacle_layer.pkl")
            if os.path.exists(path):
                try:
                    with open(path, "rb") as f:
                        loaded_layer = pickle.load(f)
                    if loaded_layer.shape == (self.map_height, self.map_width):
                        self.obstacle_layer = loaded_layer
                        self.logger.info("默认障碍层加载成功")
                        return
                    else:
                        self.logger.warning("默认障碍层形状不匹配，忽略加载")
                except Exception as e:
                    self.logger.error(f"加载默认障碍层失败: {e}")
            self.obstacle_layer = np.zeros(
                (self.map_height, self.map_width), dtype=np.uint8
            )
            self.logger.info("使用空默认障碍层")
        else:
            path = os.path.join(
                root, f"obstacle_layer{self.current_obstacle_layer_index + 1}.pkl"
            )
            if os.path.exists(path):
                try:
                    with open(path, "rb") as f:
                        loaded_layer = pickle.load(f)
                    if loaded_layer.shape == (self.map_height, self.map_width):
                        self.obstacle_layer = loaded_layer
                        self.logger.info(
                            f"预设障碍层 {self.current_obstacle_layer_index + 1} 加载成功"
                        )
                        return
                    else:
                        self.logger.warning("预设障碍层形状不匹配，忽略加载")
                except Exception as e:
                    self.logger.error(f"加载预设障碍层失败: {e}")
            else:
                self.logger.warning(f"预设文件不存在: {path}")
            self.obstacle_layer = np.zeros(
                (self.map_height, self.map_width), dtype=np.uint8
            )

    def load_next_preset(self):
        """加载下一个预设障碍层"""
        if self.obstacle_layer_file_num == 0:
            self.logger.info("无可用预设")
            return False
        self.current_obstacle_layer_index = (
            self.current_obstacle_layer_index + 1
        ) % self.obstacle_layer_file_num
        self.load_current_obstacle_layer()
        return True

    def save_current_preset(self):
        """保存当前障碍层为新预设"""
        root = os.path.join(os.path.dirname(__file__), self.map_dir_path)
        self.obstacle_layer_file_num += 1
        self.current_obstacle_layer_index = self.obstacle_layer_file_num - 1
        path = os.path.join(
            root, f"obstacle_layer{self.current_obstacle_layer_index + 1}.pkl"
        )
        try:
            with open(path, "wb") as f:
                pickle.dump(self.obstacle_layer, f)
            self.logger.info(
                f"预设障碍层 {self.current_obstacle_layer_index + 1} 保存成功"
            )
        except Exception as e:
            self.logger.error(f"保存预设障碍层失败: {e}")

    def load_map(self, yaml_path=None):
        """加载地图 yaml 和 pgm"""
        if yaml_path is None:
            yaml_path = os.path.abspath(
                os.path.join(
                    os.path.dirname(__file__), self.map_dir_path, "outdoor_final.yaml"
                )
            )
        if os.path.exists(yaml_path):
            try:
                with open(yaml_path, "r") as f:
                    map_data = yaml.safe_load(f)
                image_rel = map_data["image"]
                image_path = os.path.abspath(
                    os.path.join(os.path.dirname(yaml_path), image_rel)
                )
                self.map_image = cv2.imread(image_path, cv2.IMREAD_COLOR)
                if self.map_image is not None:
                    self.map_height, self.map_width = self.map_image.shape[:2]
                    self.resolution = map_data["resolution"]
                    self.origin = map_data["origin"][:2]

                    self.logger.info(f"地图加载成功: {image_path}")
                    return True
            except Exception as e:
                self.logger.error(f"加载地图失败: {e}")
        self.logger.warning(f"地图文件不存在: {yaml_path}")
        return False

    def qpix_generate(self):
        """生成 QPixmap"""
        if self.map_image is None:
            return
        # QPixmap for UI
        h, w, c = self.map_image.shape
        qimg = QImage(self.map_image.data, w, h, w * c, QImage.Format_RGB888)
        qimg = qimg.rgbSwapped()
        self.map_pixmap = QPixmap.fromImage(qimg)

    def load_graph_data(self, root=""):
        """加载图数据"""
        if root == "":
            root = os.path.join(os.path.dirname(__file__), self.map_dir_path)
        try:
            # 邻接矩阵
            adj_path = os.path.join(root, "adj_matrix.txt")
            if os.path.exists(adj_path):
                with open(adj_path, "r") as f:
                    self.original_node_num = int(f.readline().strip())
                    self.original_adj_matrix = np.zeros(
                        (self.original_node_num, self.original_node_num), dtype=int
                    )
                    for i in range(self.original_node_num):
                        row = list(map(int, f.readline().split()))
                        self.original_adj_matrix[i] = row
            # 节点坐标
            points_path = os.path.join(root, "points.txt")
            if os.path.exists(points_path):
                self.original_node_coordinates = []
                with open(points_path, "r") as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 2:
                            x, y = float(parts[0]), float(parts[1])
                            self.original_node_coordinates.append((x, y))
                if len(self.original_node_coordinates) != self.original_node_num:
                    raise ValueError("节点数不匹配")
            # 初始化当前为原始
            self.node_coordinates = self.original_node_coordinates[:]
            self.adj_matrix = self.original_adj_matrix.copy()
            self.node_num = self.original_node_num
            self.logger.info(f"图数据加载成功: {self.node_num} 节点")
        except Exception as e:
            self.logger.error(f"加载图数据失败: {e}")

    def get_obstacle_layer(self):
        """获取当前障碍层"""
        return self.obstacle_layer

    def set_obstacle_layer(self, layer):
        """设置当前障碍层"""
        if layer.shape == (self.map_height, self.map_width):
            self.obstacle_layer = layer

    def world_to_pixel(self, x, y):
        """世界坐标到地图像素"""
        px = int((x - self.origin[0]) / self.resolution)
        py = self.map_height - int((y - self.origin[1]) / self.resolution)
        return px, py

    def pixel_to_world(self, px, py):
        """像素到世界坐标"""
        x = px * self.resolution + self.origin[0]
        y = (self.map_height - py) * self.resolution + self.origin[1]
        return x, y

    def position_to_node(self, x, y):
        """坐标到最近节点"""
        if not self.node_coordinates:
            return 0
        min_dist = float("inf")
        closest_node = 0
        for idx, (nx, ny) in enumerate(self.node_coordinates):
            d = dist((x, y), (nx, ny))
            if d < min_dist:
                min_dist = d
                closest_node = idx
        return closest_node

    def node_to_position(self, node_id):
        """节点到坐标"""
        if 0 <= node_id < len(self.node_coordinates):
            return self.node_coordinates[node_id]
        return (0.0, 0.0)

    def node_in_obstacle(self, node_id):
        """检查节点是否在障碍内"""
        if node_id >= len(self.node_coordinates) or self.obstacle_layer.size == 0:
            return False
        x, y = self.node_coordinates[node_id]
        px, py = self.world_to_pixel(x, y)
        if 0 <= px < self.map_width and 0 <= py < self.map_height:
            return self.obstacle_layer[py, px] > 0
        return False

    def get_line_pixels(self, x0, y0, x1, y1):
        """Bresenham 算法获取线段上的像素点 (不包括端点)"""
        pixels = []
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        sx = 1 if x0 < x1 else -1
        sy = 1 if y0 < y1 else -1
        err = dx - dy
        while True:
            # 跳过端点
            if (x0 == x1 and y0 == y1) or (x0 == x0 and y0 == y0):  # 初始跳过起点
                break
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x0 += sx
            if e2 < dx:
                err += dx
                y0 += sy
            if x0 == x1 and y0 == y1:
                break
            pixels.append((x0, y0))
        return pixels

    def edge_crosses_obstacle(self, node_id1, node_id2):
        """检查边是否穿越障碍，使用 Bresenham 像素检查"""
        if node_id1 >= len(self.node_coordinates) or node_id2 >= len(
            self.node_coordinates
        ):
            return True
        start_x, start_y = self.node_coordinates[node_id1]
        end_x, end_y = self.node_coordinates[node_id2]
        px0, py0 = self.world_to_pixel(start_x, start_y)
        px1, py1 = self.world_to_pixel(end_x, end_y)
        if not (
            0 <= px0 < self.map_width
            and 0 <= py0 < self.map_height
            and 0 <= px1 < self.map_width
            and 0 <= py1 < self.map_height
        ):
            return True
        line_pixels = self.get_line_pixels(px0, py0, px1, py1)
        for px, py in line_pixels:
            if 0 <= px < self.map_width and 0 <= py < self.map_height:
                if self.obstacle_layer[py, px] > 0:
                    return True
        return False

    def apply_obstacles(self):
        """应用障碍：从原始图开始，剔除节点和穿越障碍的边"""
        if self.obstacle_layer.size == 0 or self.original_node_num == 0:
            # 重置到原始
            self.node_coordinates = self.original_node_coordinates[:]
            self.adj_matrix = self.original_adj_matrix.copy()
            self.node_num = self.original_node_num
            return
        # 从原始开始
        current_nodes = list(range(self.original_node_num))
        current_coords = self.original_node_coordinates[:]
        current_adj = self.original_adj_matrix.copy()
        current_num = self.original_node_num

        # 剔除在障碍中的节点
        valid_nodes = [
            i
            for i in current_nodes
            if not self.node_in_obstacle_with_coords(i, current_coords)
        ]
        if len(valid_nodes) == 0:
            self.logger.warning("所有节点被障碍覆盖，无法更新图")
            return

        # 更新坐标
        new_coords = [current_coords[i] for i in valid_nodes]
        # 更新邻接矩阵，检查边
        new_n = len(valid_nodes)
        new_adj = np.zeros((new_n, new_n), dtype=int)
        node_map = {old: new for new, old in enumerate(valid_nodes)}
        for i in valid_nodes:
            new_i = node_map[i]
            for j in valid_nodes:
                new_j = node_map[j]
                if current_adj[i, j] > 0:
                    # 检查边是否穿越障碍
                    if not self.edge_crosses_obstacle_with_coords(i, j, current_coords):
                        new_adj[new_i, new_j] = current_adj[i, j]

        # 更新当前
        self.node_coordinates = new_coords
        self.adj_matrix = new_adj
        self.node_num = new_n
        self.logger.info(f"障碍应用成功: 保留 {new_n} 个节点")

    def node_in_obstacle_with_coords(self, node_id, coords):
        """检查节点是否在障碍内，使用指定坐标"""
        if node_id >= len(coords) or self.obstacle_layer.size == 0:
            return False
        x, y = coords[node_id]
        px, py = self.world_to_pixel(x, y)
        if 0 <= px < self.map_width and 0 <= py < self.map_height:
            return self.obstacle_layer[py, px] > 0
        return False

    def edge_crosses_obstacle_with_coords(self, node_id1, node_id2, coords):
        """检查边是否穿越障碍，使用指定坐标"""
        if node_id1 >= len(coords) or node_id2 >= len(coords):
            return True
        start_x, start_y = coords[node_id1]
        end_x, end_y = coords[node_id2]
        px0, py0 = self.world_to_pixel(start_x, start_y)
        px1, py1 = self.world_to_pixel(end_x, end_y)
        if not (
            0 <= px0 < self.map_width
            and 0 <= py0 < self.map_height
            and 0 <= px1 < self.map_width
            and 0 <= py1 < self.map_height
        ):
            return True
        line_pixels = self.get_line_pixels(px0, py0, px1, py1)
        for px, py in line_pixels:
            if 0 <= px < self.map_width and 0 <= py < self.map_height:
                if self.obstacle_layer[py, px] > 0:
                    return True
        return False

    def save_graph_data(self, root=""):
        """保存图数据 (当前图)"""
        if root == "":
            root = os.path.join(os.path.dirname(__file__), self.map_dir_path)
        try:
            adj_path = os.path.join(root, "adj_matrix.txt")
            with open(adj_path, "w") as f:
                f.write(f"{self.node_num}\n")
                for row in self.adj_matrix:
                    f.write(" ".join(map(str, row)) + "\n")
            points_path = os.path.join(root, "points.txt")
            with open(points_path, "w") as f:
                for x, y in self.node_coordinates:
                    f.write(f"{x} {y}\n")
            self.logger.info("图数据保存成功")
        except Exception as e:
            self.logger.error(f"保存图数据失败: {e}")

    def get_obstacle_pixmap(self):
        """获取带障碍的 pixmap（绿色叠加）"""
        if self.obstacle_layer.size == 0 or self.map_image is None:
            return self.map_pixmap
        combined = self.map_image.copy()
        obstacle_mask = self.obstacle_layer > 0
        combined[obstacle_mask] = [0, 255, 0]  # BGR 绿色
        h, w, c = combined.shape
        qimg = QImage(combined.data, w, h, w * c, QImage.Format_RGB888)
        qimg = qimg.rgbSwapped()
        return QPixmap.fromImage(qimg)

    def show_editor(self):
        """显示编辑器"""
        from .obstacle_editor import ObstacleEditor

        editor = ObstacleEditor(self)
        if editor.exec_() == editor.Accepted:
            # 只更新实例，不保存图文件
            return True
        return False
